import pymysql

# 连接数据库
def connect_db(db):
    return pymysql.connect(host="localhost", port=3306, user="root", password="cxd@949846.com",database=db,charset="utf8")

# 注册功能
def register():
    con = connect_db("db_python")
    cursor = None
    try:
        cursor = con.cursor()
        username = input('请输入用户名：')
        cursor.execute(f"select * from user where username = '{username}'")
        result = cursor.fetchone()
        if result:
            print('账号已存在，注册失败！')
        else:
            password = input('请输入密码：')
            cursor.execute(f"insert into user(username,password) values('{username}','{password}')")
            con.commit()
            print('注册成功')
    except Exception as e:
        print(e)
    finally:
        if cursor:
            cursor.close()
        if con:
            con.close()

# 登录功能
def login():
    con = connect_db("db_python")
    cursor = None
    try:
        cursor = con.cursor()
        username = input('请输入用户名：')
        cursor.execute(f"select * from user where username = '{username}'")
        result = cursor.fetchone()
        if result:
            password = input('请输入密码：')
            if password == result[2]:
                print('登录成功！')
            else:
                print('密码错误，登录失败！')
        else:
            print('账户不存在')
    except Exception as e:
        print(f'发生错误{e}')
    finally:
        if cursor:
            cursor.close()
        if con:
            con.close()

# 主函数
def main():
    manu()
    while True:
        choice = input('请选择想要的操作：')
        if choice == '1':
            register()
        elif choice == '2':
            login()
        elif choice == '0':
            print('退出系统')
            break
        else:
            print('输入有误，请重新输入！')


# 菜单
def manu():
    print('************************')
    print('**** 欢迎来到洛克王国！****')
    print('****     0、退出     ****')
    print('****     1、注册     ****')
    print('****     2、登录     ****')
    print('************************')

if __name__ == "__main__":
    main()